'''
Created on Sep 10, 2010

@author: yev
'''
import array
from mpi4py import MPI
import World
import xml.etree.ElementTree

#def work(rank, name = MPI.Get_processor_name(), size, xml_file):
#    print ("I am Mario on %s" %(name))

class Assignment:
    def __init__(self, task, weight, single):
        self.task = task
        self.weight = weight
        self.single = single

def str2bool(v):
    if v is None:
        return False
    return v.lower() in ["yes", "true", "t", "1"]


def get_assignment_weights(xml_file):
    tree = xml.etree.ElementTree.ElementTree()
    fh = open(xml_file)
    root = tree.parse(fh)
    
    assignments = []
    for element in tree.getiterator('task'):
        task = element.get('id')
        weight = element.get('weight')
        single = str2bool(element.get('single'))
        assignment = Assignment(task, weight, single)
        assignments.append(assignment)
    return assignments

def who_is_who(size, comm):
    who = {}
    index = 1
    while index < size:
        where = ''
        who[index] = comm.recv(where, index, tag=World.WHO, status=None)
        print ("Mario thinks that rank %d is on %s" %(index, who[index]))
        index += 1
    return who
    
def work(xml_file):
    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    assignments = get_assignment_weights(xml_file)
    slots_available = size - 1
    
    # How many are remaining?
    remaining_assignment = 0
    for assignment in assignments:
        if assignment.weight != 'remaining':
            print assignment.weight
            slots_available =  slots_available - int(assignment.weight)
        else:
            remaining_assignment += 1

    if remaining_assignment > 1:
        print ("Error: Only one task may have the special assignment 'remaining', %d were found" %(remaining_assignment))
        who = who_is_who(size, comm)
        terminate_pipedream(who.iterkeys(), comm)
        return 
    
    if slots_available < 0:
        print ("Error: More tasks were given then number of available computing nodes. Minimum number or processors = weight_sum + 1")
        who = who_is_who(size, comm)
        terminate_pipedream(who.iterkeys(), comm)
        return 
 
    if slots_available == 0 and remaining_assignment > 0:
        print ("Error: More tasks were given then number of available computing nodes. Minimum number or processors = weight_sum + 1")
        print ("Specifically, there's no room for the task given the 'remaining' assignment'")
        who = who_is_who(size, comm)
        terminate_pipedream(who.iterkeys(), comm)
        return
    
    if slots_available > 0 and remaining_assignment == 0:
        print ("Error: You have asked for more nodes than you need. Please follow this rule: weight_sum + 1")
        who = who_is_who(size, comm)
        terminate_pipedream(who.iterkeys(), comm)
        return
     
  
  
    
    
                    
    
    for assignment in assignments:
        if assignment.weight == 'remaining':
            assignment.weight = slots_available
            print ("%d available slots were assigned to task id %d" %(slots_available, int(assignment.task)))
 
    workers = []
    index = 1
    work_assignments = {}
    who = who_is_who(size, comm)
    for a in assignments:
        print a.weight
        number_of_assignemnts = a.weight
        is_single = a.single
        hosts = []
        print ("Number of assignments for %s is %d" %(a.task, int(number_of_assignemnts)))
        number_assigned = 0;
        while number_assigned < int(number_of_assignemnts):
            print ("Mario is processing index %d" %(index))
            current_host = who[index]
            if is_single and current_host in hosts:
                print("Rank %d is being assigned %d" %(index, -1))
                comm.send(None, dest=index, tag=World.ASSIGNMENT)

            else:
                work_assignments[index] = a.task
                workers.append(index)
                print("Rank %d is being assigned %d" %(index, int(a.task)))
            
            hosts.append(current_host)
            index += 1
            number_assigned += 1
    print work_assignments
    
    index = 1;
    for index in workers:
        comm.send(work_assignments, dest=index, tag=World.ASSIGNMENT)

def terminate_pipedream(hosts, comm):
    """
    Kills all hosts
    """
    for host in hosts:
        comm.send(None, dest=host, tag=World.ASSIGNMENT) # None object causes exit

